#!/usr/bin/env python
"""
rankAbundance.py

Description:-

Give one or more pickle files as input and this shall output a white space
delimited table showing calculated biodiversity information. This includes
Shannon, SImpson & JackKnife Indices, which are indices which can be
interpreted as representations of species evenness, richness, and sampling
bias, repsectively.

Some useful options included with this program:-

    -collapse-at-rank <rank>   :-
            Collapse the result tree at rank <rank>.
    -groupX [group1files] [...] [-group2 ... [...]] :-
            * used to combine and average diversity indices over
            multiple results files. 


"""

from ssummo.ssummolib import dict_walk, Options, collapse_at_rank, combine_dicts
from math import log,fsum
import cPickle as pickle
import sys

class LocalOptions( Options ):
    def __init__( self, args ):
        Options.__init__(self)
        self.options = {'-in': [],
                           '-out' : sys.stdout,
                           '-collapse-at-rank' : None,
                           }
        self.singleargs = ['-out','-collapse-at-rank']
        self.multiargs = ['-in']
        self.helpTxt = {'-in' : 'Input pickle file. Required.',
                           '-out': 'Output table. [Standard output]',
                           '-collapse-at-rank' : 'Desired rank at which to collapse the results tree',
                          }
        self.useage = 'python {0} [options] <results file>.pkl'.format( __file__ )
        self.parseArgs( args )

class Results( object ):
    """Object for storing replicates.
    Pre-defined attributes and their indended meaning:-
        file_name   - dataset name
        xvals       - number of random samples (sequences)
        rarefaction - number of observed OTUs (OTU defined by rank)
        simpson     - simpson, as defined by Indices.Simpson()
        shannon     - shannon, as returned by Indices.Shannon()[0]
        shannon_max - shannon_max as returned by Indices.Shannon()[1]

        jack        - Once all the replicates have been calculated,
                        Results.Jack() can be called, which will
                        jackknife the replicates.
        bootstrap   - to do, not implemented yet.
    """
    __slots__ = [ 'file_name' , 'xvals' , 'rarefaction' , 'simpson' , 'shannon' ,'shannon_max' , 'jack' ]
    def __init__(self):
        """Class to store results.
        """
        #self.order = int()
        self.file_name = str()
        self.xvals = []
        self.rarefaction = {}
        self.simpson = {}
        self.shannon = {}
        self.shannon_max = {}
        self.jack = {}
    def __get__(self,val):
        return getattr( self , val )
    def __set__( self, key, val ):
        if type( val ) == Replicate:
            self.replicates[key] = val
        else:
            raise AttributeError( "{0} must be a Replicate instance".format( val.__name__ ) )
    def __getitem__(self,key):
        rep = Replicate()
        rep.nseqs = self.xvals[key]
        rep.rarefaction = self.rarefaction[key]
        rep.simpson = self.simpson[key]
        rep.shannon = self.shannon[key]
        rep.shannon_max = self.shannon_max[key]
        rep.jack = self.jack[key]
        return rep
    def __iter__(self):
        for i in xrange( len( self.xvals ) ):
            yield self[i]
    def append( self, replicate ):
        if type( replicate ) != Replicate:
            raise AttributeError( "replicate must be a Replicate instance" )
        xval = replicate.nseqs
        if xval not in self.xvals:
            self.xvals.append( xval )
            self.rarefaction.update( { xval : [] } )
            self.simpson.update( { xval : [] } )
            self.shannon.update( { xval : [] } )
            self.shannon_max.update( { xval : [] } )
            self.jack.update( { xval : [] } )
        self.rarefaction[xval].append( replicate.rarefaction )
        self.simpson[xval].append( replicate.simpson )
        self.shannon[xval].append( replicate.shannon )
        self.shannon_max[xval].append( replicate.shannon_max )
        self.jack[xval].append( replicate.jack )

    def JackKnife( self ):
        n = len( self.nseqs )
        i = 0
        theta = sum( self.jack )
        thetas = []
        while i < n:
            theta_i = theta - self.jack[i]
            theta_hat = n*theta - ( n - 1 )*theta_i
            thetas.append( theta_hat )
            i += 1
        first_order = 1/n * sum( thetas )
        self.jack = first_order
        return first_order

class Replicate( object ):
    __slots__ = [ 'nseqs', 'rarefaction' ,'shannon' , 'shannon_max' , 'simpson' ]
    def __init__(self):
        self.nseqs = int()
        for slot in self.__slots__:
            setattr( self , slot , float() )
    def __add__(self , other ):
        new = Replicate()
        for slot in self.__slots__:
            setattr( new, slot , getattr( self, slot ) + getattr( other , slot ) )
        return new
    def __div__(self,other):
        new = Replicate()
        for slot in self.__slots__:
            setattr( new , slot , getattr( self , slot ) / getattr( other , slot ) )
        return new

class Indices:
    def __init__( self, options=None ):
        if options is None:
            self.options = LocalOptions([])
        else:
            self.options = options
        self.shannon = None
        self.jack = None
        self.simpson = None
        self.name = None
    def print_header( self ):
        print '#{0} {1} {2} {3} {4}'.format( ' File'.ljust(25),
                                             'Shannon'.ljust(20) ,
                                             'Max Evenness'.ljust(25),
                                             'Simpsons'.ljust(25),
                                             'JackKnife'.ljust(25),
                                             )
        print '#{0} {1} {2} {3} {4}'.format( ' name'.ljust(25),
                                             '(H`)'.ljust(20) ,
                                             '(Hmax)'.ljust(25),
                                             'Richness'.ljust(25),
                                             '(sampling bias)'.ljust(25),
                                         )
    def __str__( self ):
        """If you haven't already calculated everything and haven't
        specified options['-out'] as anything, then we'll have to
        recalculate everything and reprint the headers for good 
        measure..."""
        if ( self.shannon and self.jack and self.simpson ) is None:
            self.calc_all()
        #    if self.options['-out'] == sys.stdout:
        #        self.__init__( self.options ) 
        if self.name is None:
            self.name = '<unknown>'
        return '{0}  {1} {2} {3} {4}'.format(   self.name,
                                      '{0:.5f}'.format( self.shannon[0] ).ljust(20) ,
                                      '{0:.5f}'.format( self.shannon[1] ).ljust(25) ,
                                      '{0:.5f}'.format( self.simpson    ).ljust(25) ,
                                      '{0:.5f}'.format( self.jack       ).ljust(25) )
    def calc_all(self):
        """Automates all calculations so that we can print the output.
        """
        counts, probs,total = self.tally( )
        richness, max_evenness = self.ShannonIndex( probs )
        diversity = self.SimpsonsIndex( counts, total )
        jack = self.jackKnife( counts, total ) 

    def tally( self , combined=False ):
        """Writes rank abundance information to self.options['-out'].
        This is represented as just a single column, showing the number
        of sequences allocated to each taxon.

        *** To use this, first define self.dictionary, preferably
                     with load( results_dict , [options] ) ***

        The function returns 3 items:-
            1) counts-a sorted list, containing the number of sequences
                      allocated to each taxa, probabilities and total number of
                      taxa.
            2) probs - a list of floating point numbers, showing the percentage
                      of sequences assigned to each taxon.
            3) total - the total number of sequences found in the dataset.
        
        These 3 items are needed to calculate the other indices... 
        """
        counts = []
        for path, node in dict_walk('', self.dictionary ):
            if 'accessions' in node:
                if combined:
                    count = 0
                    for dataset in node['accessions']:
                        for prob in dataset:
                            if float( prob ) > 0.:
                                count += 1
                    counts.append( count )
                elif len( node['accessions'] ) == 0:
                    continue
                else:
                    counts.append( len( node['accessions'] ) )
        if self.options['-out'] != sys.stdout and type(self.options['-out']) is not file:
            self.options['-out'] = file( self.options['-out'] ,'w' )
        counts = sorted(counts, reverse=True )
        total = sum(counts )

        probs = [ float(count) / total for count in counts ]
        self.options['-out'].write( '\n'.join( [ str(100.*prob) for prob in probs ] ))
        self.options['-out'].flush()
        return counts, probs, total

    def ShannonIndex( self, probs ):
        """Give the probabilities as returned from tally() and total
        species count. This returns the tuple H_prime, H_max
        """
        H_prime = - fsum( [ p * log(p) for p in probs ] )
        H_max = log( len(probs) )
        self.shannon = ( H_prime, H_max ,)
        return self.shannon

    def SimpsonsIndex( self, counts, total ):
        """Simpson's index of Diversity.
        Returns a value from 0 - 1, where 0 shows no
        diversity, and 1 shows infinite diversity.
        """
        diversity = 1. - ( fsum( [ n * (n-1.) for n in counts ] )  / (total * (total - 1. ))   )
        self.simpson = diversity
        return diversity

    def jackKnife( self, counts, total ):
        k = counts.count(1)  # number of 'unique' species.
        n = len(counts)      # number of 'plots' - observed taxa
        S = n + ( (n-1.)/n)**k  # jackknife estimate.
        self.jack = S
        return S

    def SPlusJackKnife( self, counts, total ):
        """
        Implementation of JackKnife formula from:-
        http://www.jstatsoft.org/v15/i03/paper
        Jackknife Estimator of Species Richness with 
                    S-PLUS, 2006, 15 (3)
        Smith, CD. Pontius, JS.

        J_n(S) = S0 + (n-1)/n * sum( [count for count in counts ])
        """
        S0 = len(counts)
        n = len( counts )
        JnS = total + (n-1.)/n * sum( counts )
        self.jack = JnS
        print 'S0: ',S0
        print 'n: ',n
        print 'Jns: ',JnS
        return JnS


    def load( self , results_dictionary , combined=False ,name = None):
        """Load the dictionary. If the option '-collapse-at-rank' is given,
        then we shall also collapse the dictionary at the specificied rank.
        Results dictionary is not returned, but is stored as:-
        self.dictionary
        """
        if self.options['-collapse-at-rank'] is not None:
            results_dictionary = collapse_at_rank( results_dictionary , self.options , combined=combined )
        self.dictionary = results_dictionary
        if name is not None:
            self.name = name

def main( options ):
    """Wrapper function to print everything!
    If you want to use it, just initiate a copy
    of LocalOptions, edit it's namespace variables
    as necessary (i.e. the options), and pass to
    this function.
    e.g.
    >>> myoptions = LocalOptions( list_of_input_files )
    >>> myoptions['-collapse-at-rank'] = 'genus'
    >>> main( myoptions )
    """
    in_files = []
    if len( options.regopts ) > 0:
        groups = [ group for group in sorted( options.regopts.keys() )]
        dictionaries = []
        for group in groups:
            results = []
            for infile in options.regopts[group]:
                with file( infile ,'rb' ) as handle:
                    results.append( pickle.load( handle ) )
            combined_dict = combine_dicts( results )
            dictionaries.append( combined_dict )
            in_files.append( 'GROUP{0}'.format( group ))
        combined = True
    else:
        combined = False
        dictionaries = []
        for in_file in options['-in']:
            with file( in_file , 'rb' ) as inFile:
                dictionary = pickle.load( inFile )
            dictionaries.append( dictionary )
            in_files.append( in_file )

    calculator = Indices( options )
    calculator.print_header()
    for i, in_file in enumerate( in_files ):
        calculator.load( dictionaries[i], combined , name = in_file )
        calculator.calc_all()
        print calculator

    if len( options['-in'] ) > 0:
        pass
    return dictionaries

if __name__ == '__main__':
    args = sys.argv[1:]
    myoptions = LocalOptions( args )
    main( myoptions )
